/***************************************************************************
 *   Copyright (C) 2007 by Warren Kenny   *
 *   warren.kenny@gmail.com   *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include "otaMySQLStorageEngine.h"

otaMySQLStorageEngine::otaMySQLStorageEngine()
{
}

bool otaMySQLStorageEngine::init( otaConfigReader &config )
{
	std::string serverPort;
	bool ok = ( 	config.getConfiguration( "Audit Server IP", m_serverIP ) &&
			config.getConfiguration( "Audit Server Port", serverPort ) &&
			config.getConfiguration( "DB Username", m_username ) &&
			config.getConfiguration( "DB Password", m_password ) &&
			config.getConfiguration( "DB Name", m_database ) &&
			config.getConfiguration( "DB Table Prefix", m_tablePrefix ) );
	if(ok)
	{
		std::istringstream buffer(serverPort);
		buffer >> m_serverPort;
	}
	
	mysql_init(&mysql);
	
	my_bool reconnect = 1;
	mysql_options(&mysql, MYSQL_OPT_RECONNECT, &reconnect);
	
	if(!mysql_real_connect(&mysql, m_serverIP.c_str(), m_username.c_str(), m_password.c_str(), m_database.c_str(), m_serverPort, NULL, 0))
	{
		std::cout << "Unable to connect to the MySQL Server: " << mysql_error(&mysql) << std::endl;
		ok = false;
	}
	return ok;
}
	
bool otaMySQLStorageEngine::createPreparedStatements()
{
	bool ok = true;
	
	std::string sessionDataQuery = "INSERT INTO " + m_tablePrefix + "session ( serverIP, serverPort, clientIP, clientPort, username, serverThreadID ) ";
	sessionDataQuery += "VALUES ( ?, ?, ?, ?, ?, ? )";
			
	std::string clientCommandQuery = "INSERT INTO " + m_tablePrefix + "query ( session, sequence, begin, end, exitcode, rowsReturned, bytesRec, bytesSent, packetsRec, packetsSent, queryType, db, rule ) ";
	clientCommandQuery += "VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? )";
			
	std::string clientQueryText = "INSERT INTO " + m_tablePrefix + "querytext ( session, sequence, query ) ";
	clientQueryText += "VALUES ( ?, ?, ? )";
			
	m_sessionDataQuery = mysql_stmt_init( &mysql );
	m_clientCommandQuery = mysql_stmt_init( &mysql );
	m_clientQueryText = mysql_stmt_init( &mysql );
					
	if(!m_sessionDataQuery || !m_clientCommandQuery || !m_clientQueryText)
		ok = false;
	
	if( mysql_stmt_prepare( m_sessionDataQuery, sessionDataQuery.c_str(), strlen( sessionDataQuery.c_str() ) ) != 0 )
	{
		std::cout << "Error encountered while setting up prepared statement: " << mysql_error( &mysql ) << std::endl;
		ok = false;
	}
			
	if( mysql_stmt_prepare( m_clientCommandQuery, clientCommandQuery.c_str(), strlen( clientCommandQuery.c_str() ) ) != 0 )
	{
		std::cout << "Error encountered while setting up prepared statement: " << mysql_error( &mysql ) << std::endl;
		ok = false;
	}
			
	if( mysql_stmt_prepare( m_clientQueryText, clientQueryText.c_str(), strlen( clientQueryText.c_str() ) ) != 0 )
	{
		std::cout << "Error encountered while setting up prepared statement: " << mysql_error( &mysql ) << std::endl;
		ok = false;
	}
			
	if(ok)
	{
		memset(m_sessionDataBind, 0, sizeof(m_sessionDataBind));
				
		m_sessionDataBind[0].buffer_type = MYSQL_TYPE_LONG; 
		m_sessionDataBind[0].buffer= &m_bindServerAddress;
		m_sessionDataBind[0].is_unsigned = 1;
		m_sessionDataBind[0].is_null = 0;
		m_sessionDataBind[0].length = 0;
				
		m_sessionDataBind[1].buffer_type = MYSQL_TYPE_SHORT; 
		m_sessionDataBind[1].buffer= &m_bindServerPort;
		m_sessionDataBind[1].is_unsigned = 1;
		m_sessionDataBind[1].is_null = 0;
		m_sessionDataBind[1].length = 0;
				
		m_sessionDataBind[2].buffer_type = MYSQL_TYPE_LONG; 
		m_sessionDataBind[2].buffer= &m_bindClientAddress;
		m_sessionDataBind[2].is_unsigned = 1;
		m_sessionDataBind[2].is_null = 0;
		m_sessionDataBind[2].length = 0;
				
		m_sessionDataBind[3].buffer_type = MYSQL_TYPE_SHORT; 
		m_sessionDataBind[3].buffer= &m_bindClientPort;
		m_sessionDataBind[3].is_unsigned = 1;
		m_sessionDataBind[3].is_null = 0;
		m_sessionDataBind[3].length = 0;
				
		m_sessionDataBind[4].buffer_type = MYSQL_TYPE_STRING; 
		m_sessionDataBind[4].buffer= (char *) m_bindUsername;
		m_sessionDataBind[4].buffer_length = OTA_MSE_MAX_STRING_SIZE;
		m_sessionDataBind[4].is_null = 0;
		m_sessionDataBind[4].length = &m_usernameLength;
				
		m_sessionDataBind[5].buffer_type = MYSQL_TYPE_LONG; 
		m_sessionDataBind[5].buffer= &m_bindThreadID;
		m_sessionDataBind[5].is_unsigned = 1;
		m_sessionDataBind[5].is_null = 0;
		m_sessionDataBind[5].length = 0;
				
		if( mysql_stmt_bind_param( m_sessionDataQuery, m_sessionDataBind ) != 0 )
		{
			std::cout << "Error encountered while attempting to initialize session data storage statement: " << mysql_stmt_error( m_sessionDataQuery ) << std::endl;
			ok = false;
		}
				
		memset(m_queryDataBind, 0, sizeof(m_queryDataBind));
				
		m_queryDataBind[0].buffer_type = MYSQL_TYPE_LONGLONG; 
		m_queryDataBind[0].buffer= &m_sessionID;
		m_queryDataBind[0].is_unsigned = 1;
		m_queryDataBind[0].is_null = 0;
		m_queryDataBind[0].length = 0;
				
		m_queryDataBind[1].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[1].buffer= &m_sequence;
		m_queryDataBind[1].is_unsigned = 1;
		m_queryDataBind[1].is_null = 0;
		m_queryDataBind[1].length = 0;
				
		m_queryDataBind[2].buffer_type = MYSQL_TYPE_DATETIME; 
		m_queryDataBind[2].buffer= (char *) &m_queryBegin;
		m_queryDataBind[2].is_unsigned = 1;
		m_queryDataBind[2].is_null = 0;
		m_queryDataBind[2].length = 0;
				
		m_queryDataBind[3].buffer_type = MYSQL_TYPE_DATETIME; 
		m_queryDataBind[3].buffer= (char *) &m_queryEnd;
		m_queryDataBind[3].is_unsigned = 1;
		m_queryDataBind[3].is_null = 0;
		m_queryDataBind[3].length = 0;
				
		m_queryDataBind[4].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[4].buffer= &m_exitCode;
		m_queryDataBind[4].is_null = 0;
		m_queryDataBind[4].length = 0;
				
		m_queryDataBind[5].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[5].buffer= &m_rowsReturned;
		m_queryDataBind[5].is_unsigned = 1;
		m_queryDataBind[5].is_null = 0;
		m_queryDataBind[5].length = 0;
				
		m_queryDataBind[6].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[6].buffer= &m_bytesReceived;
		m_queryDataBind[6].is_unsigned = 1;
		m_queryDataBind[6].is_null = 0;
		m_queryDataBind[6].length = 0;
				
		m_queryDataBind[7].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[7].buffer= &m_bytesSent;
		m_queryDataBind[7].is_unsigned = 1;
		m_queryDataBind[7].is_null = 0;
		m_queryDataBind[7].length = 0;
				
		m_queryDataBind[8].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[8].buffer= &m_packetsReceived;
		m_queryDataBind[8].is_unsigned = 1;
		m_queryDataBind[8].is_null = 0;
		m_queryDataBind[8].length = 0;
				
		m_queryDataBind[9].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[9].buffer= &m_packetsSent;
		m_queryDataBind[9].is_unsigned = 1;
		m_queryDataBind[9].is_null = 0;
		m_queryDataBind[9].length = 0;
				
		m_queryDataBind[10].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[10].buffer= &m_queryType;
		m_queryDataBind[10].is_unsigned = 1;
		m_queryDataBind[10].is_null = 0;
		m_queryDataBind[10].length = 0;
				
		m_queryDataBind[11].buffer_type = MYSQL_TYPE_STRING; 
		m_queryDataBind[11].buffer= (char *) m_bindDatabase;
		m_queryDataBind[11].buffer_length = OTA_MSE_MAX_STRING_SIZE;
		m_queryDataBind[11].is_null = 0;
		m_queryDataBind[11].length = &m_databaseLength;
				
		m_queryDataBind[12].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[12].buffer= &m_rule;
		m_queryDataBind[12].is_unsigned = 1;
		m_queryDataBind[12].is_null = 0;
		m_queryDataBind[12].length = 0;
								
		if( mysql_stmt_bind_param( m_clientCommandQuery, m_queryDataBind ) != 0 )
		{
			std::cout << "Error encountered while attempting to initialize query data storage statement: " << mysql_stmt_error( m_clientCommandQuery ) << std::endl;
			ok = false;
		}
				
		memset(m_queryTextBind, 0, sizeof(m_queryTextBind));
				
		m_queryTextBind[0].buffer_type = MYSQL_TYPE_LONGLONG; 
		m_queryTextBind[0].buffer= &m_sessionID;
		m_queryTextBind[0].is_unsigned = 1;
		m_queryTextBind[0].is_null = 0;
		m_queryTextBind[0].length = 0;
					
		m_queryDataBind[1].buffer_type = MYSQL_TYPE_LONG; 
		m_queryDataBind[1].buffer= &m_sequence;
		m_queryDataBind[1].is_unsigned = 1;
		m_queryDataBind[1].is_null = 0;
		m_queryDataBind[1].length = 0;
				
		m_queryTextBind[2].buffer_type = MYSQL_TYPE_STRING; 
		m_queryTextBind[2].buffer= (char *) m_bindQueryText;
		m_queryTextBind[2].buffer_length = OTA_MSE_MAX_QUERY_LENGTH;
		m_queryTextBind[2].is_null = 0;
		m_queryTextBind[2].length = &m_queryTextLength;
				
		if( mysql_stmt_bind_param( m_clientQueryText, m_queryTextBind ) != 0 )
		{
			std::cout << "Error encountered while attempting to initialize query text storage statement: " << mysql_stmt_error( m_clientQueryText ) << std::endl;
			ok = false;
		}
	}
	return ok;
}
	
bool otaMySQLStorageEngine::reset()
{
	bool ok = true;
	bool error = true;
	std::string query;
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "error";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "query";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "session";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "ruleset";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "rulemetadata";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "field";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "querytext";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "flag";
	mysql_query( &mysql, query.c_str() );
	
	query = "DROP TABLE IF EXISTS " + m_tablePrefix + "logic";
	mysql_query( &mysql, query.c_str() );
	
	query = "CREATE TABLE " + m_tablePrefix + "field ( ";
	query += "field_id int(10) NOT NULL auto_increment, ";
	query += "description varchar(64) character set utf8 NOT NULL, ";
	query += "PRIMARY KEY  (field_id) ";
	query += ") ENGINE=MyISAM DEFAULT CHARSET=latin1";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
	
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "querytext ( ";
	query += "session bigint(20) unsigned NOT NULL, ";
	query += "sequence int(10) unsigned NOT NULL, ";
	query += "query text NOT NULL ";
	query += ") ENGINE=MyISAM DEFAULT CHARSET=latin1";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
	
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "flag ( ";
	query += "flag_id int(10) NOT NULL auto_increment, ";
	query += "description varchar(64) character set utf8 NOT NULL, ";
	query += "PRIMARY KEY  (flag_id) ";
	query += ") ENGINE=MyISAM DEFAULT CHARSET=latin1";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
		
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "logic ( ";
	query += "logic_id int(10) NOT NULL auto_increment, ";
	query += "description varchar(64) character set utf8 NOT NULL, ";
	query += "PRIMARY KEY  (logic_id) ";
	query += ") ENGINE=MyISAM DEFAULT CHARSET=latin1" ;
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "ruleset ( ";
	query += "rule int(11) unsigned NOT NULL, ";
	query += "rule_meta_data_id int(10) NOT NULL, ";
	query += "field_id int(10) NOT NULL, ";
	query += "flag_id int(10) NOT NULL, ";
	query += "logic_id varchar(10) NOT NULL, ";
	query += "PRIMARY KEY  (rule) ";
	query += ") ENGINE=InnoDB DEFAULT CHARSET=utf8 ";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "rulemetadata ( ";
	query += "rule_meta_data_id int(10) unsigned NOT NULL auto_increment, ";
	query += "description varchar(64) default NULL, ";
	query += "flag int(10) NOT NULL, ";
	query += "PRIMARY KEY  (rule_meta_data_id) ";
	query += ") ENGINE=InnoDB DEFAULT CHARSET=utf8";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "session ( ";
	query += "session bigint(20) unsigned NOT NULL auto_increment, ";
	query += "serverIP int(10) unsigned NOT NULL, ";
	query += "serverPort smallint(5) unsigned NOT NULL, ";
	query += "clientIP int(10) unsigned NOT NULL, ";
	query += "clientPort smallint(5) unsigned NOT NULL, ";
	query += "username varchar(64) NOT NULL, ";
	query += "serverThreadID int(10) unsigned NOT NULL, ";
	query += "PRIMARY KEY  (session) ";
	query += ") ENGINE=InnoDB DEFAULT CHARSET=utf8";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "query ( ";
	query += "session bigint(20) unsigned NOT NULL, ";
	query += "sequence int(10) unsigned NOT NULL, ";
	query += "begin datetime default NULL, ";
	query += "end datetime default NULL, ";
	query += "exitCode int(10) unsigned NOT NULL default '0', ";
	query += "rowsReturned int(10) unsigned NOT NULL, ";
	query += "bytesRec int(10) unsigned NOT NULL, ";
	query += "bytesSent int(10) unsigned NOT NULL, ";
	query += "packetsRec int(10) unsigned NOT NULL, ";
	query += "packetsSent int(10) unsigned NOT NULL, ";
	query += "queryType tinyint(3) unsigned NOT NULL, ";
	query += "db varchar(64) NOT NULL default '', ";
	query += "rule int(10) unsigned NOT NULL, ";
	query += "PRIMARY KEY  (session, sequence), ";
	query += "KEY rule (rule), ";
	query += "CONSTRAINT Query_ibfk_1 FOREIGN KEY (session) REFERENCES " + m_tablePrefix + "session (session), ";
	query += "CONSTRAINT Query_ibfk_2 FOREIGN KEY (rule) REFERENCES " + m_tablePrefix + "ruleset (rule) ";
	query += ") ENGINE=InnoDB DEFAULT CHARSET=utf8";
	
	error = ( mysql_query( &mysql, query.c_str() ) == 0 );
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	query = "CREATE TABLE " + m_tablePrefix + "error ( ";
	query += "session bigint(20) unsigned NOT NULL, ";
	query += "sequence int(10) unsigned NOT NULL, ";
	query += "errorCode smallint(5) unsigned NOT NULL, ";
	query += "errorSeverity tinyint(3) unsigned NOT NULL, ";
	query += "errorMessage tinytext, ";
	query += "KEY session (session, sequence), ";
	query += "CONSTRAINT Error_ibfk_1 ";
	query += "FOREIGN KEY (session, sequence) ";
	query += "REFERENCES " + m_tablePrefix + "query (session, sequence) ";
	query += " ) ENGINE=InnoDB DEFAULT CHARSET=utf8 ";
		
	error = (mysql_query( &mysql, query.c_str() ) == 0);
	if(!error)
		std::cout << "Error: " << mysql_error( &mysql ) << std::endl;
			
	ok = ok && error;
	
	return ok;
}
	
void otaMySQLStorageEngine::storeSessionInformation( otaSession &session )
{
	if(!session.getSessionID( m_sessionID ))
	{
		m_bindServerAddress = session.getServerIP();
		m_bindServerPort = session.getServerPort();
		m_bindClientAddress = session.getClientIP();
		m_bindClientPort = session.getClientPort();
		m_bindThreadID = session.getThreadID();
		strcpy( m_bindUsername, session.getUsername().c_str() );
		m_usernameLength = strlen(m_bindUsername);
		
		if( mysql_stmt_execute(m_sessionDataQuery) != 0 )
		{
			if( mysql_stmt_errno( m_sessionDataQuery ) == 2013  )
			{
				bool doContinue = createPreparedStatements();
				if(!doContinue)
				{
					std::cout << "Error encountered while attempting to execute the session data storage prepared statement: " << mysql_stmt_error( m_sessionDataQuery ) << std::endl;
					return;
				}
			}
			else
			{
				std::cout << "Error encountered while attempting to execute the session data storage prepared statement: " << mysql_stmt_error( m_sessionDataQuery ) << std::endl;
				return;
			}
		}
		else
		{
			m_sessionID = mysql_stmt_insert_id( m_sessionDataQuery );
			session.setSessionID( m_sessionID );
		}	
	}
	
	otaClientCommand *command;
	m_sequence = 0;
	
	std::cout << "Attempting to close session " << m_sessionID << std::endl;
	
	while( session.getCommand( m_sequence, &command ) )
	{
		std::cout << "Storing Sequence " << m_sequence << std::endl;
		
		m_queryType = command->getType();
		
		std::string queryText;
		std::string database;
		
		queryText = command->getCommand();
		if(queryText.length() == 0)
			queryText = "No Query Text Available";	
		
		m_queryTextLength = queryText.length();
		strcpy( m_bindQueryText, queryText.c_str() );	
		
		std::cout << "Query Text Retrieved: " << m_bindQueryText << std::endl;
		
		m_queryBegin = command->getStart();
		m_queryEnd = command->getEnd();
		
		m_exitCode = command->getExitCode();
		m_rowsReturned = command->getRows();
		
		m_bytesReceived = command->getReceviedByteCount();
		m_bytesSent = command->getSentByteCount();
		
		m_packetsReceived = command->getReceivedPacketCount();
		m_packetsSent = command->getSentPacketCount();
		
		database = command->getDatabase();
		if(database.length() == 0)
			database = "Unable to identify target database";
		
		m_databaseLength = database.length() ;
		strcpy( m_bindDatabase, database.c_str() );

		std::cout << "Database Name Retrieved: " << database << std::endl;		
		
		std::cout << "Session ID = " << m_sessionID << std::endl;
		
		m_rule = 0;
		
		if( mysql_stmt_execute(m_clientCommandQuery) != 0 )
		{
			if( mysql_stmt_errno( m_clientCommandQuery ) == 2013  )
			{
				bool doContinue = createPreparedStatements();
				if(!doContinue)
				{
					std::cout << "Error encountered while attempting to execute the client command data prepared statement: " << mysql_stmt_error( m_clientCommandQuery ) << std::endl;
					return;
				}
			}
			else
			{
				std::cout << "Error encountered while attempting to execute the client command data prepared statement: " << mysql_stmt_error( m_clientCommandQuery ) << std::endl;
				return;
			}
		}
		
		if( mysql_stmt_execute(m_clientQueryText) != 0 )
		{
			if( mysql_stmt_errno( m_clientQueryText ) == 2013  )
			{
				bool doContinue = createPreparedStatements();
				if(!doContinue)
				{
					std::cout << "Error encountered while attempting to execute the client query text prepared statement: " << mysql_stmt_error( m_clientQueryText ) << std::endl;
					return;
				}
			}
			else
			{
				std::cout << "Error encountered while attempting to execute the client query text prepared statement: " << mysql_stmt_error( m_clientQueryText ) << std::endl;
				return;
			}
		}
		m_sequence++;
	}	
}

bool otaMySQLStorageEngine::close()
{
	mysql_close(&mysql);
	return true;
}


otaMySQLStorageEngine::~otaMySQLStorageEngine()
{
}


